"""
代码生成模型模块
用于集成和管理代码生成AI模型
"""

from typing import Optional, Union
from langchain_openai import OpenAI
from langchain_community.llms import Ollama
from pydantic import SecretStr
from src.config import research_settings


def get_code_generation_model(model_type: str = "openai") -> Union[OpenAI, Ollama]:
    """
    获取代码生成模型实例
    
    Args:
        model_type: 模型类型 ("openai", "ollama", "codellama")
        
    Returns:
        代码生成模型实例
    """
    if model_type == "openai":
        return OpenAI(
            base_url=research_settings.OPENAI_API_BASE,
            api_key=SecretStr(research_settings.OPENAI_API_KEY),
            model="gpt-3.5-turbo-instruct"  # 适合代码生成的模型
        )
    elif model_type == "ollama":
        return Ollama(model="codellama")
    elif model_type == "codellama":
        return Ollama(model="codellama:7b")
    else:
        raise ValueError(f"不支持的模型类型: {model_type}")